import numpy as np 
from matplotlib import pyplot as plt
import deepxde as dde
import triangle as tr
import meshio

'''------------------------------------------- Residual-based adaptive refinement ---------------------------------------------'''
def get_top_k_indices(arr, k=100):
    '''get indices of the top k values in arr'''
    arr = arr.reshape(-1)  # reshape the array to 1D
    indices = np.argsort(arr)[::-1]  # sort indices in descending order
    return indices[:k]

def plot_naca001065_add_points(points, color="k", marker="."):
    airfoil_plot = np.load('/home/aistudio/work/NACA0010-65/level1/0_Mesh_files/ordered_airfoil_points.npy')
    figure = plt.figure(dpi=200)
    axis = figure.add_subplot(111)
    axis.scatter(points[:, 0], points[:, 1], s=0.5, color=color, marker=marker)
    plt.axis("scaled")  # 设置x轴和y轴相同的缩放比例  
    plt.fill(airfoil_plot[:, 0], airfoil_plot[:, 1], color = 'white')
    plt.rcParams['ytick.direction'] = 'out'
    plt.rcParams['ytick.direction'] = 'in'
    plt.show()

def plot_naca633418_add_points(points, color="k", marker="."):
    airfoil_plot = np.load('/home/aistudio/work/NACA6334-18/level1/0_Mesh_files/ordered_airfoil_points.npy')
    figure = plt.figure(dpi=200)
    axis = figure.add_subplot(111)
    axis.scatter(points[:, 0], points[:, 1], s=0.5, color=color, marker=marker)
    plt.axis("scaled")  # 设置x轴和y轴相同的缩放比例  
    plt.fill(airfoil_plot[:, 0], airfoil_plot[:, 1], color = 'white')
    plt.rcParams['ytick.direction'] = 'out'
    plt.rcParams['ytick.direction'] = 'in'
    plt.show()

def insert_points_between_points(points):
    num_points = len(points)
    new_points = np.zeros((2*num_points, 2))
    for i in range(num_points-1):
        new_points[2*i] = points[i]
        new_points[2*i+1] = (points[i] + points[i+1]) / 2  
    new_points[-2] = points[-1]
    new_points[-1] = (points[0] + points[-1]) / 2
    return new_points

def insert_midpoints(airfoil_midpoints, airfoil_points):
    new_points = []
    for i in range(len(airfoil_points)-1):
        new_points.append(airfoil_points[i])
        midpoints_between_points = np.where(np.all(airfoil_midpoints == (airfoil_points[i] + airfoil_points[i+1]) / 2, axis=1))
        new_points.extend(airfoil_midpoints[midpoints_between_points])
    new_points.append(airfoil_points[-1])
    new_points.extend(airfoil_midpoints[np.where(np.all(airfoil_midpoints == (airfoil_points[0] + airfoil_points[-1]) / 2, axis=1))])
    return np.array(new_points)

def naca001065_geom():
    """------------------------------ NACA001065 shape -----------------------"""
    airfoil1_boundary_points = np.load('/home/aistudio/work/NACA0010-65/level1/0_Mesh_files/ordered_airfoil_points.npy')
    airfoil1 = dde.geometry.Polygon(airfoil1_boundary_points)
    """------------------------------------ Training Spatio-temporal region --------------------------------------"""
    """Spatial region definition"""
    farfield = dde.geometry.Rectangle([-5.5, -2], [6.5, 2])
    geom1 = dde.geometry.CSGDifference(farfield, airfoil1)
    return geom1

def naca633418_geom():
    """------------------------------ NACA633418 shape -----------------------"""
    airfoil2_boundary_points = np.load('/home/aistudio/work/NACA6334-18/level1/0_Mesh_files/ordered_airfoil_points.npy')
    airfoil2 = dde.geometry.Polygon(airfoil2_boundary_points)
    """------------------------------------ Training Spatio-temporal region --------------------------------------"""
    """Spatial region definition"""
    farfield = dde.geometry.Rectangle([-5.5, -2], [6.5, 2])
    geom2 = dde.geometry.CSGDifference(farfield, airfoil2)
    return geom2

"""------------------------------ Navier Stokes Equation(Time-dependent PDE) -----------------------"""
def DL_Euler_Equation_2D(x, y):
    """
    System of PDEs to be minimized: compressible 2D Euler equations.

    """
    gamma = 1.4
    
    u = y[:, 0:1]
    v = y[:, 1:2]
    p = y[:, 2:3]
    rho = y[:, 3:4]

    E = p/(gamma-1) + 0.5*rho*(u*u + v*v)

    A1, B1 = rho*u, rho*v
    A2, B2 = gamma*rho*u*u + p, gamma*rho*u*v
    A3, B3 = gamma*rho*u*v, gamma*rho*v*v + p
    A4, B4 = u*E, v*E

    dA1_x = dde.grad.jacobian(A1, x, i=0, j=0)
    dB1_y = dde.grad.jacobian(B1, x, i=0, j=1)

    dA2_x = dde.grad.jacobian(A2, x, i=0, j=0)
    dB2_y = dde.grad.jacobian(B2, x, i=0, j=1)

    dA3_x = dde.grad.jacobian(A3, x, i=0, j=0)
    dB3_y = dde.grad.jacobian(B3, x, i=0, j=1)

    dA4_x = dde.grad.jacobian(A4, x, i=0, j=0)
    dB4_y = dde.grad.jacobian(B4, x, i=0, j=1)
    
    continuity = dA1_x + dB1_y
    x_momentum = dA2_x + dB2_y
    y_momentum = dA3_x + dB3_y
    energy = dA4_x + dB4_y
    return [continuity, x_momentum, y_momentum, energy]